import streamlit as st
import os
import tempfile
import base64
import time
from PIL import Image
import io
from pic_generator import process_images_with_boots
from streamlit_paste_button import paste_image_button as pbutton

# 设置页面配置
st.set_page_config(
    page_title="AI图片生成器",
    page_icon="🎨",
    layout="wide"
)

# 添加自定义CSS样式
st.markdown("""
<style>
/* 主容器样式 */
.main-container {
    padding: 20px;
    border-radius: 10px;
    transition: all 0.3s ease;
}

/* 上传区域样式 - 独立容器悬停效果 */
.upload-section {
    background: linear-gradient(135deg, #f5f7fa 0%, #c3cfe2 100%);
    border: 3px solid transparent;
    border-radius: 20px;
    padding: 25px;
    margin: 15px 0;
    transition: all 0.4s ease;
    position: relative;
    overflow: hidden;
    box-shadow: 0 4px 15px rgba(0,0,0,0.1);
}

.upload-section:hover {
    border: 3px solid #4CAF50;
    box-shadow: 0 12px 35px rgba(76, 175, 80, 0.4);
    transform: translateY(-3px);
    background: linear-gradient(135deg, #e8f5e8 0%, #c8e6c9 100%);
}

.upload-section::before {
    content: '📤 上传区域';
    position: absolute;
    top: -15px;
    left: 20px;
    background: linear-gradient(135deg, #4CAF50 0%, #66BB6A 100%);
    color: white;
    padding: 5px 15px;
    border-radius: 15px;
    font-size: 12px;
    font-weight: bold;
    opacity: 0;
    transform: translateY(-10px);
    transition: all 0.4s ease;
    box-shadow: 0 2px 8px rgba(76, 175, 80, 0.3);
}

.upload-section:hover::before {
    opacity: 1;
    transform: translateY(0);
}

.upload-section::after {
    content: '';
    position: absolute;
    top: 0;
    left: -100%;
    width: 100%;
    height: 100%;
    background: linear-gradient(90deg, transparent, rgba(255,255,255,0.4), transparent);
    transition: left 0.6s ease;
}

.upload-section:hover::after {
    left: 100%;
}

/* 人物图片上传区域 - 作为上传区域的子元素，不独立悬停 */
.person-upload {
    background: linear-gradient(135deg, #ffecd2 0%, #fcb69f 100%);
    border: 1px solid rgba(255, 152, 0, 0.2);
    border-radius: 12px;
    padding: 15px;
    margin: 10px 0;
    transition: all 0.3s ease;
    box-shadow: 0 2px 8px rgba(255, 152, 0, 0.15);
    cursor: pointer;
    position: relative;
}

/* 服装图片上传区域 - 作为上传区域的子元素，不独立悬停 */
.clothing-upload {
    background: linear-gradient(135deg, #a8edea 0%, #fed6e3 100%);
    border: 1px solid rgba(33, 150, 243, 0.2);
    border-radius: 12px;
    padding: 15px;
    margin: 10px 0;
    transition: all 0.3s ease;
    box-shadow: 0 2px 8px rgba(33, 150, 243, 0.15);
    cursor: pointer;
    position: relative;
}

/* PROMPT区域样式 - 独立容器悬停效果 */
.prompt-section {
    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
    border: 3px solid transparent;
    border-radius: 15px;
    padding: 20px;
    margin: 10px 0;
    transition: all 0.4s ease;
    color: white;
    box-shadow: 0 4px 15px rgba(102, 126, 234, 0.3);
    position: relative;
    overflow: hidden;
}

.prompt-section:hover {
    border: 3px solid #9C27B0;
    box-shadow: 0 12px 35px rgba(156, 39, 176, 0.4);
    transform: translateY(-3px);
    background: linear-gradient(135deg, #7b1fa2 0%, #4a148c 100%);
}

.prompt-section::before {
    content: '✨ PROMPT区域';
    position: absolute;
    top: -15px;
    left: 20px;
    background: linear-gradient(135deg, #9C27B0 0%, #E91E63 100%);
    color: white;
    padding: 5px 15px;
    border-radius: 15px;
    font-size: 12px;
    font-weight: bold;
    opacity: 0;
    transform: translateY(-10px);
    transition: all 0.4s ease;
    box-shadow: 0 2px 8px rgba(156, 39, 176, 0.3);
}

.prompt-section:hover::before {
    opacity: 1;
    transform: translateY(0);
}

.prompt-section::after {
    content: '';
    position: absolute;
    top: 0;
    left: -100%;
    width: 100%;
    height: 100%;
    background: linear-gradient(90deg, transparent, rgba(255,255,255,0.2), transparent);
    transition: left 0.6s ease;
}

.prompt-section:hover::after {
    left: 100%;
}

/* 生成区域样式 - 独立容器悬停效果 */
.generation-section {
    background: linear-gradient(135deg, #ff9a9e 0%, #fecfef 100%);
    border: 3px solid transparent;
    border-radius: 15px;
    padding: 20px;
    margin: 10px 0;
    transition: all 0.4s ease;
    min-height: 400px;
    box-shadow: 0 4px 15px rgba(255, 154, 158, 0.3);
    position: relative;
    overflow: hidden;
}

.generation-section:hover {
    border: 3px solid #E91E63;
    box-shadow: 0 12px 35px rgba(233, 30, 99, 0.4);
    transform: translateY(-3px);
    background: linear-gradient(135deg, #fce4ec 0%, #f8bbd9 100%);
}

.generation-section::before {
    content: '🎨 图片生成区域';
    position: absolute;
    top: -15px;
    left: 20px;
    background: linear-gradient(135deg, #E91E63 0%, #FF5722 100%);
    color: white;
    padding: 5px 15px;
    border-radius: 15px;
    font-size: 12px;
    font-weight: bold;
    opacity: 0;
    transform: translateY(-10px);
    transition: all 0.4s ease;
    box-shadow: 0 2px 8px rgba(233, 30, 99, 0.3);
}

.generation-section:hover::before {
    opacity: 1;
    transform: translateY(0);
}

.generation-section::after {
    content: '';
    position: absolute;
    top: 0;
    left: -100%;
    width: 100%;
    height: 100%;
    background: linear-gradient(90deg, transparent, rgba(255,255,255,0.3), transparent);
    transition: left 0.6s ease;
}

.generation-section:hover::after {
    left: 100%;
}

/* 按钮悬停效果 */
.stButton > button {
    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
    color: white;
    border: none;
    border-radius: 25px;
    padding: 12px 30px;
    font-weight: bold;
    transition: all 0.3s ease;
    box-shadow: 0 4px 15px rgba(0,0,0,0.2);
    position: relative;
    overflow: hidden;
}

.stButton > button:hover {
    background: linear-gradient(135deg, #764ba2 0%, #667eea 100%);
    transform: translateY(-2px) scale(1.05);
    box-shadow: 0 6px 20px rgba(0,0,0,0.3);
}

.stButton > button::before {
    content: '';
    position: absolute;
    top: 50%;
    left: 50%;
    width: 0;
    height: 0;
    background: rgba(255,255,255,0.3);
    border-radius: 50%;
    transform: translate(-50%, -50%);
    transition: width 0.3s ease, height 0.3s ease;
}

.stButton > button:hover::before {
    width: 300px;
    height: 300px;
}

/* 文件上传器样式 */
.stFileUploader > div {
    border: 2px dashed #ccc;
    border-radius: 10px;
    padding: 20px;
    transition: all 0.3s ease;
    background: rgba(255,255,255,0.8);
    position: relative;
}

.stFileUploader > div:hover {
    border-color: #4CAF50;
    background-color: rgba(76, 175, 80, 0.1);
    transform: scale(1.02);
    box-shadow: 0 4px 12px rgba(76, 175, 80, 0.2);
}

.stFileUploader > div::after {
    content: '📁';
    position: absolute;
    top: 10px;
    right: 15px;
    font-size: 18px;
    opacity: 0.5;
    transition: all 0.3s ease;
}

.stFileUploader > div:hover::after {
    opacity: 1;
    transform: scale(1.2);
}

/* 图片预览区域 */
.image-preview {
    border: 2px solid transparent;
    border-radius: 10px;
    padding: 10px;
    margin: 5px 0;
    transition: all 0.3s ease;
    background: rgba(255,255,255,0.8);
    position: relative;
    overflow: hidden;
}

.image-preview:hover {
    border: 2px solid #FF5722;
    box-shadow: 0 4px 15px rgba(255, 87, 34, 0.3);
    transform: scale(1.02);
    background: rgba(255,255,255,0.95);
}

.image-preview::before {
    content: '';
    position: absolute;
    top: 0;
    left: 0;
    right: 0;
    bottom: 0;
    background: linear-gradient(45deg, transparent 30%, rgba(255,255,255,0.3) 50%, transparent 70%);
    transform: translateX(-100%);
    transition: transform 0.6s ease;
}

.image-preview:hover::before {
    transform: translateX(100%);
}

/* 进度条样式 */
.stProgress > div > div {
    background: linear-gradient(90deg, #4CAF50, #8BC34A, #CDDC39);
    border-radius: 10px;
    box-shadow: 0 2px 8px rgba(76, 175, 80, 0.3);
    animation: pulse 2s infinite;
}

@keyframes pulse {
    0% { box-shadow: 0 2px 8px rgba(76, 175, 80, 0.3); }
    50% { box-shadow: 0 4px 16px rgba(76, 175, 80, 0.5); }
    100% { box-shadow: 0 2px 8px rgba(76, 175, 80, 0.3); }
}

/* 成功消息样式 */
.stSuccess {
    background: linear-gradient(135deg, #c8e6c9 0%, #a5d6a7 100%);
    border-left: 5px solid #4CAF50;
    border-radius: 5px;
    animation: slideInRight 0.5s ease;
}

/* 错误消息样式 */
.stError {
    background: linear-gradient(135deg, #ffcdd2 0%, #ef9a9a 100%);
    border-left: 5px solid #f44336;
    border-radius: 5px;
    animation: shake 0.5s ease;
}

/* 信息消息样式 */
.stInfo {
    background: linear-gradient(135deg, #bbdefb 0%, #90caf9 100%);
    border-left: 5px solid #2196F3;
    border-radius: 5px;
    animation: fadeIn 0.5s ease;
}

/* 动画效果 */
@keyframes slideInRight {
    from { transform: translateX(100%); opacity: 0; }
    to { transform: translateX(0); opacity: 1; }
}

@keyframes shake {
    0%, 100% { transform: translateX(0); }
    25% { transform: translateX(-5px); }
    75% { transform: translateX(5px); }
}

@keyframes fadeIn {
    from { opacity: 0; transform: translateY(-10px); }
    to { opacity: 1; transform: translateY(0); }
}

/* 标题动画 */
h1 {
    animation: titleGlow 3s ease-in-out infinite alternate;
}

@keyframes titleGlow {
    from { text-shadow: 0 0 5px rgba(102, 126, 234, 0.5); }
    to { text-shadow: 0 0 20px rgba(102, 126, 234, 0.8), 0 0 30px rgba(118, 75, 162, 0.6); }
}

/* 响应式设计 */
@media (max-width: 768px) {
    .upload-section, .prompt-section, .generation-section {
        margin: 5px 0;
        padding: 15px;
        transform: none !important;
    }
    
    .person-upload, .clothing-upload {
        margin: 8px 0;
        padding: 12px;
    }
    
    .upload-section:hover, .prompt-section:hover, .generation-section:hover {
        transform: translateY(-1px) !important;
    }
    
    /* 移动设备上禁用子区域的独立悬停效果 */
    .upload-section:hover .person-upload,
    .upload-section:hover .clothing-upload {
        transform: none !important;
    }
}

/* 滚动条美化 */
::-webkit-scrollbar {
    width: 8px;
}

::-webkit-scrollbar-track {
    background: #f1f1f1;
    border-radius: 10px;
}

::-webkit-scrollbar-thumb {
    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
    border-radius: 10px;
}

::-webkit-scrollbar-thumb:hover {
    background: linear-gradient(135deg, #764ba2 0%, #667eea 100%);
}
</style>
""", unsafe_allow_html=True)

# 覆盖并精简的盒式样式（更贴近线框）
st.markdown("""
<style>
:root {
  --panel-bg: #ffffff;
  --panel-border: #e5e7eb;
  --panel-radius: 12px;
  --panel-shadow: 0 1px 2px rgba(0,0,0,0.06);
  --muted: #6b7280;
}

.app-wrap { padding: 8px 12px; }
.panel { background: var(--panel-bg); border: 1px solid var(--panel-border); border-radius: var(--panel-radius); padding: 16px; margin: 12px 0; box-shadow: var(--panel-shadow); }
.panel h3 { margin: 0 0 10px; font-weight: 600; color: #374151; font-size: 16px; }
.panel .hint { color: var(--muted); font-size: 12px; margin-top: 6px; }

/* 上传框更轻量样式 */
.uploader-box .stFileUploader > div { border: 1px dashed #d1d5db; border-radius: 10px; padding: 14px; background: #fafafa; }
.uploader-box .stFileUploader > div:hover { border-color: #6366f1; background: #f5f5ff; }

/* 独立上传控件容器（不被上方标题盒包裹） */
.uploader-standalone .stFileUploader > div { border: 1px dashed #d1d5db; border-radius: 10px; padding: 14px; background: #fafafa; }
.uploader-standalone .stFileUploader > div:hover { border-color: #6366f1; background: #f5f5ff; }

/* 图片预览盒 */
.preview-box { border: 1px solid var(--panel-border); border-radius: 10px; padding: 8px; margin-top: 8px; background: #fff; }

/* 生成区域面板 */
.generation-panel { min-height: 540px; }

/* 按钮更简洁样式 */
.stButton > button { background: #4f46e5; color: #fff; border: none; border-radius: 8px; padding: 10px 16px; font-weight: 600; transition: background .2s ease, transform .1s ease; }
.stButton > button:hover { background: #4338ca; transform: translateY(-1px); }
.stButton > button:active { transform: translateY(0); }
</style>
""", unsafe_allow_html=True)

def save_uploaded_file(uploaded_file):
    """保存上传的文件到临时目录"""
    try:
        with tempfile.NamedTemporaryFile(delete=False, suffix=f".{uploaded_file.name.split('.')[-1]}") as tmp_file:
            tmp_file.write(uploaded_file.getvalue())
            return tmp_file.name
    except Exception as e:
        st.error(f"保存文件失败: {e}")
        return None

def get_image_download_link(img_path, filename):
    """生成图片下载链接"""
    try:
        with open(img_path, "rb") as file:
            img_data = file.read()
        b64 = base64.b64encode(img_data).decode()
        href = f'<a href="data:image/png;base64,{b64}" download="{filename}" style="text-decoration: none; color: white; background-color: #ff4b4b; padding: 5px 10px; border-radius: 5px; font-size: 12px;">下载</a>'
        return href
    except Exception as e:
        return f"<span style='color: red;'>下载失败: {e}</span>"

def main():
    # 标题
    st.title("🎨 AI图片生成器")
    st.markdown("---")
    
    # 创建两列布局
    col1, col2 = st.columns([1, 2])
    
    with col1:
        # 人物图片上传区（独立盒）
        # 人物上传：将文件上传控件嵌入到面板盒内
        st.markdown('<div class="panel uploader-box"><div data-testid="stHeadingWithActionElements" class="panel-heading"><h3 level="3">人物图片上传区<span data-testid="stHeaderActionElements" class="panel-action"></span></h3></div><div class="hint">上传图片或拖拽</div>', unsafe_allow_html=True)
        person_images = st.file_uploader(
            "上传图片或拖拽",
            type=['png', 'jpg', 'jpeg', 'webp', 'gif'],
            accept_multiple_files=True,
            key="person_images",
            help="支持多张图片上传"
        )
        st.markdown('</div>', unsafe_allow_html=True)
        if person_images:
            st.write(f"已上传 {len(person_images)} 张人物图片")
            for i, img in enumerate(person_images[:3]):
                st.markdown('<div class="preview-box">', unsafe_allow_html=True)
                image = Image.open(img)
                st.image(image, caption=f"人物图片 {i+1}", width=150)
                st.markdown('</div>', unsafe_allow_html=True)

        # 服装上传：同样嵌入到面板盒内
        st.markdown('<div class="panel uploader-box"><div data-testid="stHeadingWithActionElements" class="panel-heading"><h3 level="3">服装图片上传区<span data-testid="stHeaderActionElements" class="panel-action"></span></h3></div><div class="hint">上传图片或拖拽</div>', unsafe_allow_html=True)
        clothing_images = st.file_uploader(
            "上传图片或拖拽",
            type=['png', 'jpg', 'jpeg', 'webp', 'gif'],
            accept_multiple_files=True,
            key="clothing_images",
            help="支持多张图片上传"
        )
        st.markdown('</div>', unsafe_allow_html=True)
        if clothing_images:
            st.write(f"已上传 {len(clothing_images)} 张服装图片")
            for i, img in enumerate(clothing_images[:3]):
                st.markdown('<div class="preview-box">', unsafe_allow_html=True)
                image = Image.open(img)
                st.image(image, caption=f"服装图片 {i+1}", width=150)
                st.markdown('</div>', unsafe_allow_html=True)

        # PROMPT 与生成（独立盒）
        st.markdown('<div class="panel">\n<h3>PROMPT区</h3>', unsafe_allow_html=True)
        custom_prompt = st.text_area(
            "文本输入区域",
            value="生成一张图片： 照片一中女生穿上牛仔裤和图二中靴子, 保持女生的姿势不变, 面部特征表情、背景和整体真实感不变。广角镜头保证女生全身和靴子均在图内。",
            height=100,
            help="输入生成图片的提示词"
        )
        generate_button = st.button("点击生成", type="primary", use_container_width=True)
        st.markdown('</div>', unsafe_allow_html=True)
    
    with col2:

        
        # 初始化session state
        if 'generated_images' not in st.session_state:
            st.session_state.generated_images = []
        
        # 生成图片逻辑
        if generate_button:
            if not person_images or not clothing_images:
                st.error("请先上传人物图片和服装图片！")
            else:
                # 清空之前的结果和显示区域
                st.session_state.generated_images = []
                
                # 计算总组合数并显示预估时间
                total_combinations = len(person_images) * len(clothing_images)
                estimated_time = total_combinations * 3 + (total_combinations - 1) * 3  # 每张图3秒生成 + 间隔3秒
                st.info(f"将生成 {total_combinations} 张图片，预计需要 {estimated_time // 60} 分 {estimated_time % 60} 秒")
                
                with st.spinner("正在生成图片，请稍候..."):
                    progress_bar = st.progress(0)
                    current_progress = 0
                    success_count = 0
                    failed_count = 0
                    
                    for i, person_img in enumerate(person_images):
                        for j, clothing_img in enumerate(clothing_images):
                            try:
                                # 添加延迟避免API调用过于频繁
                                if current_progress > 0:  # 第一次调用不延迟
                                    st.info(f"等待 3 秒后继续生成下一张图片...")
                                    time.sleep(3)  # 每次API调用间隔3秒
                                
                                # 保存上传的文件到临时目录
                                person_path = save_uploaded_file(person_img)
                                clothing_path = save_uploaded_file(clothing_img)
                                
                                if person_path and clothing_path:
                                    st.info(f"正在生成第 {current_progress + 1}/{total_combinations} 张图片...")
                                    
                                    # 调用图片生成函数
                                    result_path = process_images_with_boots(
                                        person_path, 
                                        clothing_path, 
                                        custom_prompt
                                    )
                                    
                                    if result_path and os.path.exists(result_path):
                                        st.session_state.generated_images.append({
                                            'path': result_path,
                                            'person_name': person_img.name,
                                            'clothing_name': clothing_img.name,
                                            'combination': f"人物{i+1}_服装{j+1}"
                                        })
                                        success_count += 1
                                        st.success(f"第 {current_progress + 1} 张图片生成成功！")
                                    else:
                                        failed_count += 1
                                        st.warning(f"第 {current_progress + 1} 张图片生成失败，跳过...")
                                    
                                    # 清理临时文件
                                    try:
                                        os.unlink(person_path)
                                        os.unlink(clothing_path)
                                    except:
                                        pass
                                
                                current_progress += 1
                                progress_bar.progress(current_progress / total_combinations)
                                
                            except Exception as e:
                                failed_count += 1
                                st.error(f"生成图片失败 (人物{i+1}_服装{j+1}): {e}")
                                st.write(f"详细错误信息: {str(e)}")
                                current_progress += 1
                                progress_bar.progress(current_progress / total_combinations)
                    
                    progress_bar.empty()
                    
                    # 显示最终统计
                    if success_count > 0:
                        st.success(f"批量生成完成！成功: {success_count} 张，失败: {failed_count} 张")
                    else:
                        st.error(f"批量生成失败！所有 {total_combinations} 张图片都生成失败")
                        st.info("建议检查网络连接或稍后重试")
        
        # 显示生成的图片（多宫格布局）
        if st.session_state.generated_images:
            st.success(f"成功生成 {len(st.session_state.generated_images)} 张图片！")
            
            # 计算网格布局
            num_images = len(st.session_state.generated_images)
            cols_per_row = 3  # 每行3张图片
            
            for i in range(0, num_images, cols_per_row):
                cols = st.columns(cols_per_row)
                
                for j in range(cols_per_row):
                    if i + j < num_images:
                        img_data = st.session_state.generated_images[i + j]
                        
                        with cols[j]:
                            try:
                                # 显示图片
                                image = Image.open(img_data['path'])
                                st.markdown('<div class="image-preview">', unsafe_allow_html=True)
                                st.image(image, caption=img_data['combination'], use_column_width=True)
                                st.markdown('</div>', unsafe_allow_html=True)
                                
                                # 下载按钮
                                download_filename = f"{img_data['combination']}_generated.png"
                                download_link = get_image_download_link(img_data['path'], download_filename)
                                st.markdown(download_link, unsafe_allow_html=True)
                                
                            except Exception as e:
                                st.error(f"显示图片失败: {e}")
        else:
            # 显示占位符
            st.info("👆 请上传图片并点击生成按钮开始创作")
            st.image("https://via.placeholder.com/600x400/f0f0f0/999999?text=图片生成区域", 
                    caption="等待生成图片...", use_column_width=True)
        
        st.markdown('</div>', unsafe_allow_html=True)  # 关闭生成区域容器

if __name__ == "__main__":
    main()